"""GPU配置管理"""

from enum import Enum
from typing import List, Dict, Any, Optional


class GpuType(Enum):
    """GPU类型枚举"""
    NVIDIA = "nvidia"
    AMD = "amd"
    INTEL = "intel"
    AUTO = "auto"


class GpuConfig:
    """GPU配置类"""
    
    def __init__(self):
        self.available_gpus: List[Dict[str, Any]] = []
        self.preferred_gpu: GpuType = GpuType.AUTO
        self.max_gpu_usage: int = 80  # %
        self.temperature_limit: int = 85  # °C
        self.memory_threshold: int = 90  # %
        
    def detect_gpus(self) -> List[Dict[str, Any]]:
        """检测可用GPU"""
        gpus = []
        
        try:
            # 尝试使用GPUtil检测NVIDIA GPU
            import GPUtil
            nvidia_gpus = GPUtil.getGPUs()
            for gpu in nvidia_gpus:
                gpus.append({
                    "type": GpuType.NVIDIA,
                    "name": gpu.name,
                    "id": gpu.id,
                    "memory_total": gpu.memoryTotal,
                    "memory_used": gpu.memoryUsed,
                    "temperature": gpu.temperature,
                    "load": gpu.load * 100,
                })
        except ImportError:
            pass
        
        # 检测AMD GPU (需要额外库支持)
        # 检测Intel GPU (需要额外库支持)
        
        self.available_gpus = gpus
        return gpus
    
    def get_optimal_gpu(self) -> Optional[Dict[str, Any]]:
        """获取最优GPU"""
        if not self.available_gpus:
            self.detect_gpus()
        
        if not self.available_gpus:
            return None
        
        # 根据偏好选择GPU
        if self.preferred_gpu != GpuType.AUTO:
            for gpu in self.available_gpus:
                if gpu["type"] == self.preferred_gpu:
                    return gpu
        
        # 自动选择：优先选择负载最低的GPU
        return min(self.available_gpus, key=lambda g: g.get("load", 100))
    
    def can_use_gpu(self, gpu_info: Dict[str, Any]) -> bool:
        """检查GPU是否可用"""
        if gpu_info.get("load", 0) > self.max_gpu_usage:
            return False
        if gpu_info.get("temperature", 0) > self.temperature_limit:
            return False
        if gpu_info.get("memory_used", 0) / gpu_info.get("memory_total", 1) * 100 > self.memory_threshold:
            return False
        return True
    
    def get_gpu_status(self) -> Dict[str, Any]:
        """获取GPU状态汇总"""
        if not self.available_gpus:
            self.detect_gpus()
        
        total_gpus = len(self.available_gpus)
        available_gpus = [g for g in self.available_gpus if self.can_use_gpu(g)]
        
        return {
            "total_gpus": total_gpus,
            "available_gpus": len(available_gpus),
            "gpu_list": self.available_gpus,
            "optimal_gpu": self.get_optimal_gpu(),
        }